#include <iostream>
#include <cstring>
#include <assert.h>
#include "tagger.h"
#include "util.h"
namespace MRFPP
{

void Tagger::open(FeatureIndex* index,int n,bool isTraining)
{
	feature_index_ = index;
	ysize_ = index -> ysize();
	if(isTraining)
		xsize_ = index -> xsize(n);
	else
		xsize_ = n;
	index -> buildFeatures(this);
}

void Tagger::fillNodePotential(size_t cur, size_t y, std::vector <int>& f)
{
	nodePotential[cur][y] = 0;
	for(size_t j = 0; j < f.size(); j++)
		if ( f[j] + y < alpha_.size())
		nodePotential[cur][y] +=  alpha_[f[j] + y];
}
/**
we only have one edge feature
*/
void Tagger::fillEdgePotential(size_t x,size_t y,size_t cur,std::vector <int>& f)
{
	edgePotential[cur][x][y] = 0;
	if (f.size() && (f[0] + x*ysize_ + y ) < alpha_.size())
		edgePotential[cur][x][y] = alpha_[f[0] + x*ysize_ + y];
}

void Tagger::updateNodePotential(size_t cur,std::vector <int>& f)
{
	for(size_t j = 0; j < f.size(); j++)
	{

		alpha_[f[j] + answer[cur]] ++;
		alpha_[f[j] + result[cur]] --;
	}
}

void Tagger::updateEdgePotential(size_t x,size_t cur, std::vector <int>& f)
{
	alpha_[f[0] + x*ysize_ + answer[cur]] ++;
	alpha_[f[0] + x*ysize_ + result[cur]] --;
}

int Tagger::viterbi(bool print)
{
	for(size_t i = 0; i < xsize_; i++)
	{
		if ( i == 0 )
		{
			for( size_t j = 0; j < ysize_; j++)
				bestCost[i][j] = nodePotential[i][j];
			continue;
		}

		for(size_t j = 0; j < ysize_; j++)
		{
			bestCost[i][j] = -1e37;
			for(size_t k = 0; k < ysize_; k++)
			{
				double tmpCost = bestCost[i - 1][k] + edgePotential[i][k][j];
				if ( tmpCost > bestCost[i][j])
				{
					bestCost[i][j] = tmpCost;
					pos[i][j] = k;
				}
			}

			bestCost[i][j] += nodePotential[i][j];
		}

	}
	double maxCost = 0.0;
	int maxy = 0;
	for(size_t i = 0; i < ysize_; i++)
	{
		if(bestCost[xsize_ - 1][i] > maxCost)
		{
			maxCost = bestCost[xsize_ - 1][i];
			maxy = i;
		}
	}
	int last;
	result.resize(xsize_);
	result[xsize_ - 1] = maxy;
	for(int i = xsize_ - 1; i > 0 ; i--)
	{
		maxy = pos[i][maxy];
		result[i - 1] = maxy;
	}
	int right = 0;
	for(int i = 0; i < xsize_; i++)
	{

		if(result[i] == answer[i])
			right++;
		if(print)
			std::cout << feature_index_->getTag(result[i]) << " ";
	}
	if(print)
		std::cout << std::endl;
	//std::cout << right / xsize_ << std::endl;
	return right;
}

void Tagger::forwardbackward()
{
	for(size_t j = 0; j < ysize_; j++)
	{
		forwardMatrix[0][j] = nodePotential[0][j];
	}

	for(size_t i = 1; i < xsize_; i++)
	{
		for(size_t j = 0; j < ysize_; j++)
		{
			for(size_t k = 0; k < ysize_; k++)
			{
				forwardMatrix[i][j] = logsumexp(forwardMatrix[i][j], forwardMatrix[i-1][k] + edgePotential[i][k][j],(k == 0));
			}
			forwardMatrix[i][j] += nodePotential[i][j];
		}
	}

	for(size_t j = 0; j < ysize_; j++ )
		backwardMatrix[xsize_ - 1][j] = nodePotential[xsize_ -1][j];
	for(int i = xsize_ - 2; i >= 0; --i)
	{
		for(size_t j = 0; j < ysize_; j++)
		{
			for(size_t k = 0; k < ysize_; k++)
			{
				backwardMatrix[i][j] = logsumexp(backwardMatrix[i][j], backwardMatrix[i + 1][k] + edgePotential[i+1][j][k],(k == 0));
			}
			backwardMatrix[i][j] += nodePotential[i][j];
		}
	}

	Z_ = 0;
	for(size_t i = 0; i < ysize_; i++)
		Z_ = logsumexp(Z_,backwardMatrix[0][i] , i == 0);
}
}
